{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "pd.set_option('display.float_format', '{:.2f}'.format)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>person</th>\n",
       "      <th>year</th>\n",
       "      <th>spending</th>\n",
       "      <th>inp</th>\n",
       "      <th>out</th>\n",
       "      <th>out_rand</th>\n",
       "      <th>mental</th>\n",
       "      <th>mdcl</th>\n",
       "      <th>drugs</th>\n",
       "      <th>supplies</th>\n",
       "      <th>...</th>\n",
       "      <th>rand_plan_group6</th>\n",
       "      <th>fam_size</th>\n",
       "      <th>position</th>\n",
       "      <th>first_month_site</th>\n",
       "      <th>contact_month_site</th>\n",
       "      <th>start_month</th>\n",
       "      <th>start_month_site</th>\n",
       "      <th>mde</th>\n",
       "      <th>coins</th>\n",
       "      <th>coins_dental</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>MA250247</td>\n",
       "      <td>1.00</td>\n",
       "      <td>14.35</td>\n",
       "      <td>0.00</td>\n",
       "      <td>14.35</td>\n",
       "      <td>14.35</td>\n",
       "      <td>0.00</td>\n",
       "      <td>14.35</td>\n",
       "      <td>14.35</td>\n",
       "      <td>0.00</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>4.00</td>\n",
       "      <td>3.00</td>\n",
       "      <td>181.00</td>\n",
       "      <td>173.00</td>\n",
       "      <td>181.00</td>\n",
       "      <td>181.00</td>\n",
       "      <td>1000.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>1.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>MA250247</td>\n",
       "      <td>2.00</td>\n",
       "      <td>115.15</td>\n",
       "      <td>0.00</td>\n",
       "      <td>115.15</td>\n",
       "      <td>115.15</td>\n",
       "      <td>0.00</td>\n",
       "      <td>115.15</td>\n",
       "      <td>24.65</td>\n",
       "      <td>0.00</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>4.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>181.00</td>\n",
       "      <td>173.00</td>\n",
       "      <td>193.00</td>\n",
       "      <td>193.00</td>\n",
       "      <td>1000.00</td>\n",
       "      <td>0.95</td>\n",
       "      <td>0.95</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>MA250247</td>\n",
       "      <td>4.00</td>\n",
       "      <td>110.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>110.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>4.00</td>\n",
       "      <td>4.00</td>\n",
       "      <td>181.00</td>\n",
       "      <td>173.00</td>\n",
       "      <td>217.00</td>\n",
       "      <td>217.00</td>\n",
       "      <td>1000.00</td>\n",
       "      <td>0.95</td>\n",
       "      <td>0.95</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>MA250255</td>\n",
       "      <td>2.00</td>\n",
       "      <td>75.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>75.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>4.00</td>\n",
       "      <td>2.00</td>\n",
       "      <td>181.00</td>\n",
       "      <td>173.00</td>\n",
       "      <td>193.00</td>\n",
       "      <td>193.00</td>\n",
       "      <td>1000.00</td>\n",
       "      <td>0.95</td>\n",
       "      <td>0.95</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>MA250255</td>\n",
       "      <td>3.00</td>\n",
       "      <td>15.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>15.00</td>\n",
       "      <td>15.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>15.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>4.00</td>\n",
       "      <td>4.00</td>\n",
       "      <td>181.00</td>\n",
       "      <td>173.00</td>\n",
       "      <td>205.00</td>\n",
       "      <td>205.00</td>\n",
       "      <td>1000.00</td>\n",
       "      <td>0.95</td>\n",
       "      <td>0.95</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 307 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "     person  year  spending  inp    out  out_rand  mental   mdcl  drugs  \\\n",
       "0  MA250247  1.00     14.35 0.00  14.35     14.35    0.00  14.35  14.35   \n",
       "1  MA250247  2.00    115.15 0.00 115.15    115.15    0.00 115.15  24.65   \n",
       "2  MA250247  4.00    110.00 0.00 110.00      0.00    0.00   0.00   0.00   \n",
       "3  MA250255  2.00     75.00 0.00  75.00      0.00    0.00   0.00   0.00   \n",
       "4  MA250255  3.00     15.00 0.00  15.00     15.00    0.00  15.00   0.00   \n",
       "\n",
       "   supplies  ...  rand_plan_group6  fam_size  position  first_month_site  \\\n",
       "0      0.00  ...                 0      4.00      3.00            181.00   \n",
       "1      0.00  ...                 0      4.00      1.00            181.00   \n",
       "2      0.00  ...                 0      4.00      4.00            181.00   \n",
       "3      0.00  ...                 0      4.00      2.00            181.00   \n",
       "4      0.00  ...                 0      4.00      4.00            181.00   \n",
       "\n",
       "   contact_month_site  start_month  start_month_site     mde  coins  \\\n",
       "0              173.00       181.00            181.00 1000.00   1.00   \n",
       "1              173.00       193.00            193.00 1000.00   0.95   \n",
       "2              173.00       217.00            217.00 1000.00   0.95   \n",
       "3              173.00       193.00            193.00 1000.00   0.95   \n",
       "4              173.00       205.00            205.00 1000.00   0.95   \n",
       "\n",
       "   coins_dental  \n",
       "0          1.00  \n",
       "1          0.95  \n",
       "2          0.95  \n",
       "3          0.95  \n",
       "4          0.95  \n",
       "\n",
       "[5 rows x 307 columns]"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = pd.read_stata(\n",
    "    \"../data/person_years.dta\"\n",
    ")\n",
    "\n",
    "data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2773388/441753252.py:24: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  data['plantype'] = data['plan'].astype(int).apply(group_plantypes)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>person</th>\n",
       "      <th>year</th>\n",
       "      <th>spending</th>\n",
       "      <th>inp</th>\n",
       "      <th>out</th>\n",
       "      <th>out_rand</th>\n",
       "      <th>mental</th>\n",
       "      <th>mdcl</th>\n",
       "      <th>drugs</th>\n",
       "      <th>supplies</th>\n",
       "      <th>...</th>\n",
       "      <th>fam_size</th>\n",
       "      <th>position</th>\n",
       "      <th>first_month_site</th>\n",
       "      <th>contact_month_site</th>\n",
       "      <th>start_month</th>\n",
       "      <th>start_month_site</th>\n",
       "      <th>mde</th>\n",
       "      <th>coins</th>\n",
       "      <th>coins_dental</th>\n",
       "      <th>plantype</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>MA250247</td>\n",
       "      <td>1.00</td>\n",
       "      <td>14.35</td>\n",
       "      <td>0.00</td>\n",
       "      <td>14.35</td>\n",
       "      <td>14.35</td>\n",
       "      <td>0.00</td>\n",
       "      <td>14.35</td>\n",
       "      <td>14.35</td>\n",
       "      <td>0.00</td>\n",
       "      <td>...</td>\n",
       "      <td>4.00</td>\n",
       "      <td>3.00</td>\n",
       "      <td>181.00</td>\n",
       "      <td>173.00</td>\n",
       "      <td>181.00</td>\n",
       "      <td>181.00</td>\n",
       "      <td>1000.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>MA250247</td>\n",
       "      <td>2.00</td>\n",
       "      <td>115.15</td>\n",
       "      <td>0.00</td>\n",
       "      <td>115.15</td>\n",
       "      <td>115.15</td>\n",
       "      <td>0.00</td>\n",
       "      <td>115.15</td>\n",
       "      <td>24.65</td>\n",
       "      <td>0.00</td>\n",
       "      <td>...</td>\n",
       "      <td>4.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>181.00</td>\n",
       "      <td>173.00</td>\n",
       "      <td>193.00</td>\n",
       "      <td>193.00</td>\n",
       "      <td>1000.00</td>\n",
       "      <td>0.95</td>\n",
       "      <td>0.95</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>MA250247</td>\n",
       "      <td>4.00</td>\n",
       "      <td>110.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>110.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>...</td>\n",
       "      <td>4.00</td>\n",
       "      <td>4.00</td>\n",
       "      <td>181.00</td>\n",
       "      <td>173.00</td>\n",
       "      <td>217.00</td>\n",
       "      <td>217.00</td>\n",
       "      <td>1000.00</td>\n",
       "      <td>0.95</td>\n",
       "      <td>0.95</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>MA250255</td>\n",
       "      <td>2.00</td>\n",
       "      <td>75.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>75.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>...</td>\n",
       "      <td>4.00</td>\n",
       "      <td>2.00</td>\n",
       "      <td>181.00</td>\n",
       "      <td>173.00</td>\n",
       "      <td>193.00</td>\n",
       "      <td>193.00</td>\n",
       "      <td>1000.00</td>\n",
       "      <td>0.95</td>\n",
       "      <td>0.95</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>MA250255</td>\n",
       "      <td>3.00</td>\n",
       "      <td>15.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>15.00</td>\n",
       "      <td>15.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>15.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>...</td>\n",
       "      <td>4.00</td>\n",
       "      <td>4.00</td>\n",
       "      <td>181.00</td>\n",
       "      <td>173.00</td>\n",
       "      <td>205.00</td>\n",
       "      <td>205.00</td>\n",
       "      <td>1000.00</td>\n",
       "      <td>0.95</td>\n",
       "      <td>0.95</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20198</th>\n",
       "      <td>MH320080</td>\n",
       "      <td>3.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>...</td>\n",
       "      <td>6.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>1362.00</td>\n",
       "      <td>6000.00</td>\n",
       "      <td>251.00</td>\n",
       "      <td>1506.00</td>\n",
       "      <td>154.01</td>\n",
       "      <td>0.95</td>\n",
       "      <td>0.95</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20199</th>\n",
       "      <td>MH320301</td>\n",
       "      <td>3.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>...</td>\n",
       "      <td>7.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>1374.00</td>\n",
       "      <td>6000.00</td>\n",
       "      <td>253.00</td>\n",
       "      <td>1518.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20200</th>\n",
       "      <td>MH320458</td>\n",
       "      <td>5.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>...</td>\n",
       "      <td>7.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>1218.00</td>\n",
       "      <td>6000.00</td>\n",
       "      <td>251.00</td>\n",
       "      <td>1506.00</td>\n",
       "      <td>450.00</td>\n",
       "      <td>0.95</td>\n",
       "      <td>0.95</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20201</th>\n",
       "      <td>MH320640</td>\n",
       "      <td>3.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>...</td>\n",
       "      <td>4.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>1374.00</td>\n",
       "      <td>6000.00</td>\n",
       "      <td>253.00</td>\n",
       "      <td>1518.00</td>\n",
       "      <td>524.70</td>\n",
       "      <td>0.25</td>\n",
       "      <td>0.50</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20202</th>\n",
       "      <td>MH320738</td>\n",
       "      <td>3.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>...</td>\n",
       "      <td>2.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>1362.00</td>\n",
       "      <td>6000.00</td>\n",
       "      <td>251.00</td>\n",
       "      <td>1506.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>20203 rows × 308 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "         person  year  spending  inp    out  out_rand  mental   mdcl  drugs  \\\n",
       "0      MA250247  1.00     14.35 0.00  14.35     14.35    0.00  14.35  14.35   \n",
       "1      MA250247  2.00    115.15 0.00 115.15    115.15    0.00 115.15  24.65   \n",
       "2      MA250247  4.00    110.00 0.00 110.00      0.00    0.00   0.00   0.00   \n",
       "3      MA250255  2.00     75.00 0.00  75.00      0.00    0.00   0.00   0.00   \n",
       "4      MA250255  3.00     15.00 0.00  15.00     15.00    0.00  15.00   0.00   \n",
       "...         ...   ...       ...  ...    ...       ...     ...    ...    ...   \n",
       "20198  MH320080  3.00      0.00 0.00   0.00      0.00    0.00   0.00   0.00   \n",
       "20199  MH320301  3.00      0.00 0.00   0.00      0.00    0.00   0.00   0.00   \n",
       "20200  MH320458  5.00      0.00 0.00   0.00      0.00    0.00   0.00   0.00   \n",
       "20201  MH320640  3.00      0.00 0.00   0.00      0.00    0.00   0.00   0.00   \n",
       "20202  MH320738  3.00      0.00 0.00   0.00      0.00    0.00   0.00   0.00   \n",
       "\n",
       "       supplies  ...  fam_size  position  first_month_site  \\\n",
       "0          0.00  ...      4.00      3.00            181.00   \n",
       "1          0.00  ...      4.00      1.00            181.00   \n",
       "2          0.00  ...      4.00      4.00            181.00   \n",
       "3          0.00  ...      4.00      2.00            181.00   \n",
       "4          0.00  ...      4.00      4.00            181.00   \n",
       "...         ...  ...       ...       ...               ...   \n",
       "20198      0.00  ...      6.00      1.00           1362.00   \n",
       "20199      0.00  ...      7.00      1.00           1374.00   \n",
       "20200      0.00  ...      7.00      1.00           1218.00   \n",
       "20201      0.00  ...      4.00      1.00           1374.00   \n",
       "20202      0.00  ...      2.00      1.00           1362.00   \n",
       "\n",
       "       contact_month_site  start_month  start_month_site     mde  coins  \\\n",
       "0                  173.00       181.00            181.00 1000.00   1.00   \n",
       "1                  173.00       193.00            193.00 1000.00   0.95   \n",
       "2                  173.00       217.00            217.00 1000.00   0.95   \n",
       "3                  173.00       193.00            193.00 1000.00   0.95   \n",
       "4                  173.00       205.00            205.00 1000.00   0.95   \n",
       "...                   ...          ...               ...     ...    ...   \n",
       "20198             6000.00       251.00           1506.00  154.01   0.95   \n",
       "20199             6000.00       253.00           1518.00    0.00   0.00   \n",
       "20200             6000.00       251.00           1506.00  450.00   0.95   \n",
       "20201             6000.00       253.00           1518.00  524.70   0.25   \n",
       "20202             6000.00       251.00           1506.00    0.00   0.00   \n",
       "\n",
       "       coins_dental  plantype  \n",
       "0              1.00         4  \n",
       "1              0.95         4  \n",
       "2              0.95         4  \n",
       "3              0.95         4  \n",
       "4              0.95         4  \n",
       "...             ...       ...  \n",
       "20198          0.95         4  \n",
       "20199          0.00         1  \n",
       "20200          0.95         2  \n",
       "20201          0.50         3  \n",
       "20202          0.00         1  \n",
       "\n",
       "[20203 rows x 308 columns]"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def group_plantypes(plan: int) -> int:\n",
    "    \"\"\"Group plans into four types:\n",
    "        1-Free\n",
    "        2-Individual Deductible\n",
    "        3-Cost Sharing (25%/50%)\n",
    "        4-Catostrophic (Fam Deductible) (95%/100%)\n",
    "\n",
    "    Args:\n",
    "        plan (int)\n",
    "\n",
    "    Returns:\n",
    "        int: Plantype\n",
    "    \"\"\"\n",
    "    if plan == 24:\n",
    "        return 1\n",
    "    if plan in [1, 5]:\n",
    "        return 2\n",
    "    if plan in range(9, 24):\n",
    "        return 3\n",
    "    if plan in [*range(2,5), *range(6,9)]:\n",
    "        return 4\n",
    "    raise ValueError(\"Invalid input\")\n",
    "\n",
    "data['plantype'] = data['plan'].astype(int).apply(group_plantypes)\n",
    "\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th>plantype</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>count</th>\n",
       "      <td>6840.00</td>\n",
       "      <td>4175.00</td>\n",
       "      <td>5464.00</td>\n",
       "      <td>3724.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>4.45</td>\n",
       "      <td>2.98</td>\n",
       "      <td>3.27</td>\n",
       "      <td>2.78</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>7.41</td>\n",
       "      <td>5.66</td>\n",
       "      <td>5.71</td>\n",
       "      <td>5.50</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>min</th>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25%</th>\n",
       "      <td>1.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50%</th>\n",
       "      <td>2.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>2.00</td>\n",
       "      <td>1.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>75%</th>\n",
       "      <td>5.00</td>\n",
       "      <td>3.00</td>\n",
       "      <td>4.00</td>\n",
       "      <td>3.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>max</th>\n",
       "      <td>144.00</td>\n",
       "      <td>83.00</td>\n",
       "      <td>77.00</td>\n",
       "      <td>105.00</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "plantype       1       2       3       4\n",
       "count    6840.00 4175.00 5464.00 3724.00\n",
       "mean        4.45    2.98    3.27    2.78\n",
       "std         7.41    5.66    5.71    5.50\n",
       "min         0.00    0.00    0.00    0.00\n",
       "25%         1.00    0.00    0.00    0.00\n",
       "50%         2.00    1.00    2.00    1.00\n",
       "75%         5.00    3.00    4.00    3.00\n",
       "max       144.00   83.00   77.00  105.00"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.groupby(\"plantype\")['ftf'].describe().T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mastering-py-metrics",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
